{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://github.com/ritheshkumar95/pytorch-vqvae/tree/master"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.distributions.normal import Normal\n",
    "from torch.distributions.kl import kl_divergence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VectorQuantization(Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, inputs, codebook):\n",
    "        with torch.no_grad():\n",
    "            embedding_size = codebook.size(1)\n",
    "            inputs_size = inputs.size()\n",
    "            inputs_flatten = inputs.view(-1, embedding_size)\n",
    "\n",
    "            codebook_sqr = torch.sum(codebook ** 2, dim=1)\n",
    "            inputs_sqr = torch.sum(inputs_flatten ** 2, dim=1, keepdim=True)\n",
    "\n",
    "            # Compute the distances to the codebook\n",
    "            distances = torch.addmm(codebook_sqr + inputs_sqr,\n",
    "                inputs_flatten, codebook.t(), alpha=-2.0, beta=1.0)\n",
    "\n",
    "            _, indices_flatten = torch.min(distances, dim=1)\n",
    "            indices = indices_flatten.view(*inputs_size[:-1])\n",
    "            ctx.mark_non_differentiable(indices)\n",
    "\n",
    "            return indices\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad_output):\n",
    "        raise RuntimeError('Trying to call `.grad()` on graph containing '\n",
    "            '`VectorQuantization`. The function `VectorQuantization` '\n",
    "            'is not differentiable. Use `VectorQuantizationStraightThrough` '\n",
    "            'if you want a straight-through estimator of the gradient.')\n",
    "\n",
    "class VectorQuantizationStraightThrough(Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, inputs, codebook):\n",
    "        indices = vq(inputs, codebook)\n",
    "        indices_flatten = indices.view(-1)\n",
    "        ctx.save_for_backward(indices_flatten, codebook)\n",
    "        ctx.mark_non_differentiable(indices_flatten)\n",
    "\n",
    "        codes_flatten = torch.index_select(codebook, dim=0,\n",
    "            index=indices_flatten)\n",
    "        codes = codes_flatten.view_as(inputs)\n",
    "\n",
    "        return (codes, indices_flatten)\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad_output, grad_indices):\n",
    "        grad_inputs, grad_codebook = None, None\n",
    "\n",
    "        if ctx.needs_input_grad[0]:\n",
    "            # Straight-through estimator\n",
    "            grad_inputs = grad_output.clone()\n",
    "        if ctx.needs_input_grad[1]:\n",
    "            # Gradient wrt. the codebook\n",
    "            indices, codebook = ctx.saved_tensors\n",
    "            embedding_size = codebook.size(1)\n",
    "\n",
    "            grad_output_flatten = (grad_output.contiguous()\n",
    "                                              .view(-1, embedding_size))\n",
    "            grad_codebook = torch.zeros_like(codebook)\n",
    "            grad_codebook.index_add_(0, indices, grad_output_flatten)\n",
    "\n",
    "        return (grad_inputs, grad_codebook)\n",
    "\n",
    "vq = VectorQuantization.apply\n",
    "vq_st = VectorQuantizationStraightThrough.apply\n",
    "\n",
    "\n",
    "class VQEmbedding(nn.Module):\n",
    "    def __init__(self, K, D):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(K, D)\n",
    "        self.embedding.weight.data.uniform_(-1./K, 1./K)\n",
    "        \n",
    "        self.vq = VectorQuantization()\n",
    "        self.vq_st = VectorQuantizationStraightThrough()\n",
    "\n",
    "    def forward(self, z_e_x):\n",
    "        z_e_x_ = z_e_x.permute(0, 2, 3, 1).contiguous()\n",
    "        latents = vq(z_e_x_, self.embedding.weight)\n",
    "        return latents\n",
    "\n",
    "    def straight_through(self, z_e_x):\n",
    "        z_e_x_ = z_e_x.permute(0, 2, 3, 1).contiguous()\n",
    "        z_q_x_, indices = vq_st(z_e_x_, self.embedding.weight.detach())\n",
    "        z_q_x = z_q_x_.permute(0, 3, 1, 2).contiguous()\n",
    "\n",
    "        z_q_x_bar_flatten = torch.index_select(self.embedding.weight,\n",
    "            dim=0, index=indices)\n",
    "        z_q_x_bar_ = z_q_x_bar_flatten.view_as(z_e_x_)\n",
    "        z_q_x_bar = z_q_x_bar_.permute(0, 3, 1, 2).contiguous()\n",
    "\n",
    "        return z_q_x, z_q_x_bar\n",
    "\n",
    "\n",
    "class ResBlock(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.block = nn.Sequential(\n",
    "            nn.ReLU(True),\n",
    "            nn.Conv2d(dim, dim, 3, 1, 1),\n",
    "            nn.BatchNorm2d(dim),\n",
    "            nn.ReLU(True),\n",
    "            nn.Conv2d(dim, dim, 1),\n",
    "            nn.BatchNorm2d(dim)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x + self.block(x)\n",
    "\n",
    "\n",
    "class VectorQuantizedVAE(nn.Module):\n",
    "    def __init__(self, input_dim, dim, K=512):\n",
    "        super().__init__()\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Conv2d(input_dim, dim, 4, 2, 1),\n",
    "            nn.BatchNorm2d(dim),\n",
    "            nn.ReLU(True),\n",
    "            nn.Conv2d(dim, dim, 4, 2, 1),\n",
    "            ResBlock(dim),\n",
    "            ResBlock(dim),\n",
    "        )\n",
    "\n",
    "        self.codebook = VQEmbedding(K, dim)\n",
    "\n",
    "        self.decoder = nn.Sequential(\n",
    "            ResBlock(dim),\n",
    "            ResBlock(dim),\n",
    "            nn.ReLU(True),\n",
    "            nn.ConvTranspose2d(dim, dim, 4, 2, 1),\n",
    "            nn.BatchNorm2d(dim),\n",
    "            nn.ReLU(True),\n",
    "            nn.ConvTranspose2d(dim, input_dim, 4, 2, 1),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "\n",
    "        self.apply(weights_init)\n",
    "\n",
    "    def encode(self, x):\n",
    "        z_e_x = self.encoder(x)\n",
    "        latents = self.codebook(z_e_x)\n",
    "        return latents\n",
    "\n",
    "    def decode(self, latents):\n",
    "        z_q_x = self.codebook.embedding(latents).permute(0, 3, 1, 2)  # (B, D, H, W)\n",
    "        x_tilde = self.decoder(z_q_x)\n",
    "        return x_tilde\n",
    "\n",
    "    def forward(self, x):\n",
    "        z_e_x = self.encoder(x)\n",
    "        z_q_x_st, z_q_x = self.codebook.straight_through(z_e_x)\n",
    "        x_tilde = self.decoder(z_q_x_st)\n",
    "        return x_tilde, z_e_x, z_q_x"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
